import copy
from abc import abstractmethod
from typing import Any

from torch import nn

from backbone import get_backbone
from classifier import build_continual_classifier


# abstract class
class Base(nn.Module):
    # fc: ContinualClf
    fc_func: str
    constant_dim: bool

    def __init__(self, args):
        super(Base, self).__init__()

        print("This is for the BaseNet initialization.")
        self.backbone = get_backbone(args)
        self.fc = build_continual_classifier(
            name=args["classifier"],
            in_features=self.backbone.out_dim,
            out_features=args["init_cls"],
            fc_func=args["fc_func"],
            freeze_old=args["freeze_old"],
            constant_dim=args["constant_dim"],
            **args.get("fc_kwargs", {}),
        )
        print("After BaseNet initialization.")
        self._device = args["device"][0]
        self.init_cls = args["init_cls"]
        self.inc_cls = args["inc_cls"]
        self.out_dim: int = self.backbone.out_dim

    @property
    def feature_dim(self):
        return self.out_dim

    @abstractmethod
    def forward(self, x, *args, **kwargs) -> Any:
        pass

    @abstractmethod
    def extract_feats(self, x):
        pass

    @abstractmethod
    def extract_token(self, x):
        pass

    def copy(self):
        return copy.deepcopy(self)

    def freeze(self):
        for param in self.parameters():
            param.requires_grad = False
        self.eval()

        return self

    def show_trainable_params(self):
        for name, param in self.named_parameters():
            if param.requires_grad:
                print(name, param.numel())

    def before_task(self):
        self.backbone.before_task()  # type: ignore
        self.fc.before_task()

    def after_train(self):
        self.backbone.after_train()  # type: ignore
        self.fc.after_train()

    def after_task(self):
        self.backbone.after_task()  # type: ignore
        self.fc.after_task()

    def after_medium(self):
        self.backbone.after_medium()  # type: ignore
        self.fc.after_medium()
